package ratelimit

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"
	"net/http"
	"time"
)

type Builder struct {
	//前缀
	prefix string
	//redis客户端
	cmd redis.Cmdable
	//时间间隔
	interval time.Duration
	//令牌桶容量
	rate int
}

var luaScript string

// 一个中间件是一个示例，用示例可以添加自定义信息，如果是函数自定义设置就麻烦了
// 用于构建一个Builder实例
func NewBuilder(cmd redis.Cmdable, interval time.Duration, rate int) *Builder {
	return &Builder{
		cmd:      cmd,
		prefix:   "ip-limiter",
		interval: interval,
		rate:     rate,
	}
}

// 用于构建一个中间件
func (b *Builder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		limited, err := b.limit(ctx)
		if err != nil {
			ctx.AbortWithStatus(http.StatusInternalServerError)
			return
		}
		if limited {
			ctx.AbortWithStatus(http.StatusTooManyRequests)
			return
		}
	}
}

// 用于判断是否超出限制
func (b *Builder) limit(ctx *gin.Context) (bool, error) {
	key := fmt.Sprintf("%s:%s", b.prefix, ctx.ClientIP())
	//Eval方法用于执行Lua脚本
	return b.cmd.Eval(ctx, luaScript, []string{key}, b.rate, b.interval.Milliseconds()).Bool()
}
